package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.{StringIndexer, OneHotEncoder => OH}
import org.apache.spark.sql.DataFrame

/**
  * @Author: dongguosheng, songhaicheng
  * @Date: 2018/8/16 15:40
  */
case class OneHotEncoderParam(
                               override val input_pt: String,
                               override val output_pt: String,
                               override val hive_table: String,
                               override val flow_time: String,
                               override val inputCol: String,
                               override val outputCol: String,
                               override val preserveCols: String,
                               override val modelPath: String,
                               // 如何处理无效值，支持 [skip, error, keep]，默认 error
                               handleInvalid: String,
                               // 编码全为 0 的向量标签是否可以占有单独的二进制特征，默认 false
                               saveLast: Boolean
                             ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, "error", false)

  override def verify(): Unit = {
    super.verify()
    val handles = Array("skip", "error", "keep")
    require(handles.contains(handleInvalid.toLowerCase),
      s"param handleInvalid only accepts ${handles.mkString("[", ", ", "]")}, but has $handleInvalid")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("handleInvalid" -> handleInvalid)
    map += ("saveLast" -> saveLast)
    map
  }
}


class OneHotEncoder extends AbstractFeatureEncoder[OneHotEncoderParam] {

  override def setUp(param: OneHotEncoderParam, data: DataFrame): Array[PipelineStage] = {
    val idx = new StringIndexer()
      .setInputCol(param.inputCol)
      .setOutputCol(tool.renameDuplicatedColName(s"${param.inputCol}_idx", data.schema.fieldNames))
      .setHandleInvalid(param.handleInvalid)

    val oh = new OH()
      .setInputCol(idx.getOutputCol)
      .setOutputCol(param.outputCol)
      .setDropLast(!param.saveLast)

    Array(idx, oh)
  }

}

object OneHotEncoder {

  def apply(paramStr: String): Unit = {
    new OneHotEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    OneHotEncoder(args(0))
  }
}
